package zws.cassandra3;

import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.handler.codec.LineBasedFrameDecoder;
import io.netty.handler.codec.string.LineEncoder;
import io.netty.handler.codec.string.StringDecoder;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;
import org.apache.cassandra.cql3.*;
import org.apache.cassandra.cql3.statements.*;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.Keyspace;
import org.apache.cassandra.db.Memtable;
import org.apache.cassandra.db.partitions.PartitionUpdate;
import org.apache.cassandra.exceptions.RequestExecutionException;
import org.apache.cassandra.exceptions.RequestValidationException;
import org.apache.cassandra.service.QueryState;
import org.apache.cassandra.transport.messages.ResultMessage;
import org.apache.cassandra.utils.MD5Digest;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;
import java.lang.management.ManagementFactory;
import java.nio.ByteBuffer;
import java.util.*;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

public class TestQueryHandler implements QueryHandler {

    private static final Logger logger = LoggerFactory.getLogger(TestQueryHandler.class);

    static SnapshotManager<Map<String, List<PartitionUpdate>>> cassandraSnapshotManager = new EmbeddedCassandraSnapshotManager();

    private static volatile boolean isInit = false;

    public static volatile boolean needTimeout = false;
    public static volatile String dataManipulationType;
    public static volatile String targetTable;

    private QueryHandler instance;

    public TestQueryHandler() {
        if (!isInit) {
            synchronized (TestQueryHandler.class) {
                if (!isInit) {
                    try {
                        int port = 19042;
                        // port
                        if (System.getProperties().containsKey("backdoor.port")) {
                            logger.info("backdoor.port have been passed: {}", System.getProperty("backdoor.port"));
                            port = Integer.parseInt(System.getProperty("backdoor.port"));
                        }
                        // suicide
                        if (System.getProperties().containsKey("backdoor.suicide")) {
                            logger.info("flag backdoor.suicide have been set");
                            int suicideAfterMin = 20;
                            if (System.getProperties().containsKey("backdoor.suicideAfterMin")) {
                                logger.info("backdoor.suicideAfterMin have been passed: {}", System.getProperty("backdoor.suicideAfterMin"));
                                suicideAfterMin = Integer.parseInt(System.getProperty("backdoor.suicideAfterMin"));
                            }
                            Executors.newSingleThreadScheduledExecutor().schedule(() -> {
                                    logger.info("standalone cassandra start to suicide...");
                                    System.exit(0);
                                }, suicideAfterMin, TimeUnit.MINUTES);
                        }
                        bootTestBackdoor(port);
                        printPid();
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }

                    isInit = true;
                }
            }
        }
        this.instance = QueryProcessor.instance;
    }

    @Override
    public ResultMessage process(String query, QueryState state, QueryOptions options, Map<String, ByteBuffer> customPayload, long queryStartNanoTime) throws RequestExecutionException, RequestValidationException {
        return instance.process(query, state, options, customPayload, queryStartNanoTime);
    }

    @Override
    public ResultMessage.Prepared prepare(String query, QueryState state, Map<String, ByteBuffer> customPayload) throws RequestValidationException {
        return instance.prepare(query, state, customPayload);
    }

    @Override
    public ParsedStatement.Prepared getPrepared(MD5Digest id) {
        return instance.getPrepared(id);
    }

    @Override
    public ParsedStatement.Prepared getPreparedForThrift(Integer id) {
        return instance.getPreparedForThrift(id);
    }

    @Override
    public ResultMessage processPrepared(CQLStatement statement, QueryState state, QueryOptions options, Map<String, ByteBuffer> customPayload, long queryStartNanoTime) throws RequestExecutionException, RequestValidationException {
        waitForTimeoutIfNeeded(statement);
        return instance.processPrepared(statement, state, options, customPayload, queryStartNanoTime);
    }

    @Override
    public ResultMessage processBatch(BatchStatement statement, QueryState state, BatchQueryOptions options, Map<String, ByteBuffer> customPayload, long queryStartNanoTime) throws RequestExecutionException, RequestValidationException {
        return instance.processBatch(statement, state, options, customPayload, queryStartNanoTime);
    }

    private void waitForTimeoutIfNeeded(CQLStatement statement) {
        if (!needTimeout) {
            return;
        }

        //
        if (dataManipulationType != null) {
            switch (dataManipulationType.toUpperCase()) {
                case "INSERT":
                case "UPDATE":
                    if (!(statement instanceof UpdateStatement)) {
                        return;
                    }
                    break;
                case "DELETE":
                    if (!(statement instanceof DeleteStatement)) {
                        return;
                    }
                    break;
                case "SELECT":
                    if (!(statement instanceof SelectStatement)) {
                        return;
                    }
                    break;
                default:
                    return;
            }
        }

        logger.debug("targetTable: {}", targetTable);
        if (targetTable != null) {
            if (statement instanceof ModificationStatement) {
                System.out.println("stmtTable: " + ((ModificationStatement) statement).columnFamily());
                if (!targetTable.equalsIgnoreCase(((ModificationStatement) statement).columnFamily())) {
                    return;
                }

            } else {
                if (statement instanceof SelectStatement) {
                    logger.debug("stmtTable: {}", ((SelectStatement) statement).columnFamily());
                    if (!targetTable.equalsIgnoreCase(((SelectStatement) statement).columnFamily())) {
                        return;
                    }

                } else {
                    return;
                }
            }
        }

        try {
            TimeUnit.MILLISECONDS.sleep(2500);
        } catch (Throwable t) {

        }

    }


    public static void bootTestBackdoor(int port) throws InterruptedException {
        EventLoopGroup workerGroup = new NioEventLoopGroup(1);
        new ServerBootstrap()
                .group(workerGroup, workerGroup)
                .channel(NioServerSocketChannel.class)
                .handler(new LoggingHandler(LogLevel.INFO))
                .childHandler(new TestBackdoorChannelInitializer())
                .bind(port)
                .sync();
        logger.info("test backdoor started, port: {}", port);
    }

    private static void printPid() {
        // get pid
        String processName = ManagementFactory.getRuntimeMXBean().getName();
        logger.info("pid: {}", processName.substring(0, processName.indexOf('@')));
    }

    private static class TestBackdoorChannelInitializer extends ChannelInitializer<SocketChannel> {
        @Override
        protected void initChannel(SocketChannel ch) throws Exception {
            ChannelPipeline pipeline = ch.pipeline();
            pipeline.addLast(new LineBasedFrameDecoder(Integer.MAX_VALUE));
            pipeline.addLast(new StringDecoder());
            pipeline.addLast(new LineEncoder());
//            pipeline.addLast(new LoggingHandler(LogLevel.INFO));
            pipeline.addLast(new TestBackdoorChannelInboundHandler());
        }
    }

    private static class TestBackdoorChannelInboundHandler extends SimpleChannelInboundHandler<String> {
        @Override
        protected void channelRead0(ChannelHandlerContext ctx, String message) throws Exception {
            logger.debug("backdoor recv: {}", message);
            try {
                if (StringUtils.isBlank(message)) {
                    ctx.writeAndFlush("Ok.");
                    return;
                }
                String[] array = message.split(" ");
                UnsafeCassandraOperation op = UnsafeCassandraOperation.valueOf(array[0]);
                String[] args = Arrays.copyOfRange(array, 1, array.length);

                switch (op) {
                    case TABLE_TRUNCATE:
                        truncateTables(args);
                        break;
                    case SNAPSHOT_TAKE:
                        takeSnapshot(args);
                        break;
                    case SNAPSHOT_RESTORE:
                        restoreSnapshot(args);
                        break;
                    case SNAPSHOT_DELETE:
                        deleteSnapshot(args);
                        break;
                    case SNAPSHOT_CLEAR:
                        clearSnapshots();
                        break;
                    case TIMEOUT_ENABLE:
                        enableTimeout(args);
                        break;
                    case TIMEOUT_DISABLE:
                        disableTimeout();
                        break;
                    case GC:
                        gc();
                        break;
                    default:
                        throw new UnsupportedOperationException();
                }
                ctx.writeAndFlush("Ok.");
            } catch (Exception e) {
                ByteArrayOutputStream baos = new ByteArrayOutputStream(64 * 1024);
                e.printStackTrace(new PrintStream(baos));
                ctx.writeAndFlush("Failure. " + baos.toString().replace('\n', ';'));
                e.printStackTrace();
            }
        }
    }

    private static void truncateTables(String[] args) {
        String keyspace = args[0];
        String[] tables = Arrays.copyOfRange(args, 1, args.length);
        logger.debug("truncate tables, keyspace: {}, tables: {}", keyspace, Arrays.toString(tables));
        logger.trace("on heap used before truncate: {}", Memtable.MEMORY_POOL.onHeap.used());
        for (String table : tables) {
            Keyspace ks = Keyspace.open(keyspace);
            ColumnFamilyStore cfs = ks.getColumnFamilyStore(table.trim());
            // clear counter that record memory usage
            cfs.concatWithIndexes().forEach(s -> {
                s.getTracker().getView().getAllMemtables().forEach(m -> m.getAllocator().setDiscarded());
            });
            // call the unsafe method to truncate
            cfs.clearUnsafe();

            if (Objects.nonNull(ks.viewManager.forTable(cfs.metadata))) {
                for (ColumnFamilyStore viewCfs : ks.viewManager.forTable(cfs.metadata).allViewsCfs()) {
                    viewCfs.concatWithIndexes().forEach(s -> {
                        s.getTracker().getView().getAllMemtables().forEach(m -> m.getAllocator().setDiscarded());
                    });
                    viewCfs.clearUnsafe();
                }
            }
        }
        logger.trace("on heap used after truncate: {}", Memtable.MEMORY_POOL.onHeap.used());
    }

    private static void takeSnapshot(String[] args) {
        String snapshot = args[0];
        String keyspace = args[1];

        Map<String, Object> arguments = new HashMap<>();
        arguments.put("keyspace", keyspace);
        cassandraSnapshotManager.snapshot0(snapshot, arguments);
    }

    private static void restoreSnapshot(String[] args) {
        String snapshot = args[0];
        String keyspace = args[1];

        Map<String, Object> arguments = new HashMap<>();
        arguments.put("keyspace", keyspace);
        cassandraSnapshotManager.restore0(snapshot, arguments);
    }

    private static void deleteSnapshot(String[] args) {
        String snapshot = args[0];
        cassandraSnapshotManager.delete(snapshot);
    }

    private static void clearSnapshots() {
        cassandraSnapshotManager.clear();
    }

    private static void enableTimeout(String[] args) {
        String dataManipulationType = args[0];
        String targetTable = args[1];

        TestQueryHandler.needTimeout = true;

        if ("NO_DATA_MANIPULATION_TYPE".equalsIgnoreCase(dataManipulationType)) {
            TestQueryHandler.dataManipulationType = null;
        } else {
            TestQueryHandler.dataManipulationType = dataManipulationType;
        }

        if ("NO_TARGET_TABLE".equalsIgnoreCase(targetTable)) {
            TestQueryHandler.targetTable = null;
        } else {
            TestQueryHandler.targetTable = targetTable;
        }
    }

    private static void disableTimeout() {
        TestQueryHandler.needTimeout = false;
        TestQueryHandler.dataManipulationType = null;
        TestQueryHandler.targetTable = null;
    }

    private static void gc() {
        System.gc();
    }

}
